import numpy as np
import matplotlib.pyplot as plt
from collections import Counter
import numpy.random as npr
import matplotlib
matplotlib.use(backend="TkAgg")

"""
    计算几何分布（首次成功前的失败次数）的期望。

    分析过程：
    看第一步的结果（第一试）：
    - 若第一试成功（概率 p），那么过程立刻结束，失败次数 X = 0。
    - 若第一试失败（概率 1 - p），这次失败“浪费”了 1 次，然后过程从头开始（独立同分布），剩下的失败次数分布和原来一样，期望仍为 E[X]。

    因此根据条件期望可得：
    E[X] = p * 0 + (1 - p) * (1 + E[X])

    右边解释：若失败，已经发生了 1 次失败，再加上之后仍期望为 E[X] 的失败数。

    解这个方程：
    E[X] = (1 - p)(1 + E[X]) = (1 - p) + (1 - p)E[X]

    把含 E[X] 的项移到左边：
    E[X] - (1 - p)E[X] = 1 - p ⇒ E[X](1 - (1 - p)) = 1 - p

    注意 1 - (1 - p) = p，所以：
    pE[X] = 1 - p ⇒ E[X] = (1 - p) / p

    :param p: 每次试验成功的概率
    :return: 首次成功前失败次数的期望
    """


p=0.3 # success probability on each independent trial
N=200000 # number of independent trials

Y = npr.geometric(p=p, size=N) # number of trials until success (include success)
X  =Y-1 # number of failures before success

# empirical and theoretical means
empirical_mean_X = X.mean()
theoretical_mean_X = (1 - p) / p     # E[X] for "failures before first success"
empirical_mean_Y = Y.mean()
theoretical_mean_Y = 1 / p           # E[Y] for "trials until first success"

print(f"p = {p}, simulations = {N}")
print(f"Empirical mean of X (#failures) = {empirical_mean_X:.6f}")
print(f"Theoretical E[X] = (1-p)/p = {theoretical_mean_X:.6f}")
print(f"Empirical mean of Y (#trials) = {empirical_mean_Y:.6f}")
print(f"Theoretical E[Y] = 1/p = {theoretical_mean_Y:.6f}")

# prepare histogram for X and theoretical pmf
max_k = 15 # show first several k values
counts = Counter(X)
ks = np.arange(0, max_k+1)
emp_pmf = np.array([counts[k] for k in ks]) /N
theo_pmf = (1 - p) ** ks * p


# plot
plt.figure(figsize=(9,5))
plt.bar(ks, emp_pmf, width=0.4, label='Empirical PMF (X)', alpha=0.6, align='center')
plt.plot(ks, theo_pmf, 'o-', color='red', label='Theoretical PMF (X)')
plt.xticks(ks)
plt.xlabel('k (number of failures before first success)')
plt.ylabel('Probability P(X=k)')
plt.title('Geometric distribution (failures before first success): empirical vs theoretical')
plt.legend()
plt.tight_layout()
plt.show()
